# visual_verify_original_wildrgbd.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
from tqdm import tqdm
import json
import cv2

# --- Import from our visual utility file ---
from visual_util import launch_visualizer

# --- Import from the original utility file for data loading helpers ---
try:
    from verify_utils import setup_logging, depthmap_to_absolute_camera_coordinates, imread_cv2
except ImportError:
    print("FATAL ERROR: Could not import from 'verify_utils.py'.")
    print("Please ensure 'verify_utils.py' is in the same directory.")
    sys.exit(1)


def load_original_wildrgbd_view_data(rgb_path, camera_intrinsics, poses_dict, logger, use_mask=True):
    """
    Loads all necessary data for a single view from the ORIGINAL wildrgbd dataset format.
    """
    try:
        frame_id_str = osp.basename(rgb_path).replace('.png', '')
        frame_id = int(frame_id_str)

        if frame_id not in poses_dict:
            logger.warning(f"No pose found for frame {frame_id}. Skipping.")
            return None

        # Construct paths
        scene_dir = osp.dirname(osp.dirname(rgb_path)) # up two levels from rgb/
        depth_path = osp.join(scene_dir, "depth", f"{frame_id:05d}.png")
        mask_path = osp.join(scene_dir, "masks", f"{frame_id:05d}.png")

        if not osp.exists(depth_path):
            logger.warning(f"Missing depth file for {frame_id_str}.png. Skipping.")
            return None

        # Load image and depth data
        rgb_image = imread_cv2(rgb_path) # Assumes it returns RGB
        depthmap = imread_cv2(depth_path, -1) # Use -1 for cv2.IMREAD_UNCHANGED
        print(rgb_image.shape)
        sys.exit(0)
        
        if rgb_image is None or depthmap is None:
            logger.warning(f"Failed to read image or depth for {rgb_path}. Skipping.")
            return None

        # --- Process data ---
        # 1. Normalize RGB image for the viewer
        rgb_image = rgb_image.astype(np.float32) / 255.0
        # 2. Assume original depth is in mm and scale to meters.
        depthmap = depthmap.astype(np.float32) / 1000.0
        depthmap[~np.isfinite(depthmap)] = 0
        
        # 3. Optionally apply the foreground mask
        if use_mask:
            if not osp.exists(mask_path):
                logger.warning(f"Mask file not found for {frame_id_str}.png. Proceeding without mask.")
            else:
                object_mask = imread_cv2(mask_path, 0) # Load as grayscale
                if object_mask is not None:
                    object_mask_bool = (object_mask.astype(np.float32) / 255.0) > 0.1
                    depthmap *= object_mask_bool

        # Get camera pose from the pre-loaded dictionary
        camera_pose = poses_dict[frame_id]

        return {
            "rgb": rgb_image, "depth": depthmap, "intrinsics": camera_intrinsics,
            "pose": camera_pose, "path": rgb_path
        }
    except Exception as e:
        logger.error(f"Could not process view for {osp.basename(rgb_path)}. Details: {e}", exc_info=True)
        return None

# ======================================================================================
# --- Main Execution Block ---
# ======================================================================================

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Visually verify the consistency of an ORIGINAL wildrgbd dataset."
    )
    
    parser.add_argument(
        "--data_root", 
        type=str,
        default='/lc/data/3D/wildrgbd/original/orig',
        help="Path to the root directory of the wildrgbd dataset."
    )
    
    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/visual_verification_original_wildrgbd.log",
        help="Path to the output log file."
    )

    parser.add_argument(
        "--category", 
        type=str, 
        default='apple',
        help="Specify the category to visualize (e.g., 'doll')."
    )

    parser.add_argument(
        "--scene", 
        type=str, 
        default='scene_018',
        help="Specify the scene name to visualize (e.g., 'scene0001_00')."
    )

    parser.add_argument(
        "--port", 
        type=int, 
        default=8080,
        help="Port to run the Viser web server on."
    )
    
    parser.add_argument(
        "--no_mask",
        action="store_true",
        help="If set, do not apply the foreground mask to the depth map."
    )
    
    args = parser.parse_args()
    args.no_mask = True
    
    log_dir = osp.dirname(args.log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    logger = setup_logging(args.log_file)

    scene_dir = osp.join(args.data_root, args.category, 'scenes', args.scene)
    if not osp.isdir(scene_dir):
        logger.error(f"Specified scene directory does not exist: {scene_dir}")
        sys.exit(1)

    # --- Load Metadata and Poses for the entire scene once ---
    try:
        with open(osp.join(scene_dir, "metadata"), "r") as f:
            metadata = json.load(f)
        
        K_raw = np.array(metadata["K"]).reshape(3, 3).T
        fx, fy, cx, cy = K_raw[0, 0], K_raw[1, 1], K_raw[0, 2], K_raw[1, 2]
        camera_intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)

        poses_content = np.genfromtxt(osp.join(scene_dir, "cam_poses.txt"))
        poses_dict = {
            int(row[0]): row[1:].reshape(4, 4).astype(np.float32)
            for row in poses_content
        }
    except Exception as e:
        logger.error(f"Failed to load metadata or poses for scene {args.scene}. Error: {e}")
        sys.exit(1)


    # --- Process the specified scene ---
    scene_key = f"{args.category}/{args.scene}"
    logger.info(f"\n--- Processing scene: {scene_key} ---")

    rgb_dir = osp.join(scene_dir, "rgb")
    rgb_paths = sorted(glob.glob(osp.join(rgb_dir, "*.png")))
    if not rgb_paths:
        logger.warning(f"No RGB images found in {rgb_dir}. Skipping.")
        sys.exit(1)

    pc_list, color_list, conf_list = [], [], []
    focals, pps, Rs, ts = [], [], [], []
    
    logger.info("Loading views and generating point clouds...")
    for rgb_path in tqdm(rgb_paths, file=sys.stdout):
        view_data = load_original_wildrgbd_view_data(
            rgb_path, camera_intrinsics, poses_dict, logger, use_mask=not args.no_mask
        )
        if view_data:
            world_points, valid_mask, _ = depthmap_to_absolute_camera_coordinates(
                view_data["depth"], view_data["intrinsics"], view_data["pose"]
            )
            pc_list.append(world_points)
            color_list.append(view_data['rgb'])
            conf_list.append(valid_mask.astype(np.float32))

            intrinsics = view_data['intrinsics']
            pose = view_data['pose']
            focals.append(intrinsics[0, 0])
            pps.append((intrinsics[0, 2], intrinsics[1, 2]))
            Rs.append(pose[:3, :3])
            ts.append(pose[:3, 3])

    if not pc_list:
        logger.warning(f"No valid data could be loaded for {scene_key}. Skipping visualization.")
        sys.exit(1)

    cam_dict = {'focal': focals, 'pp': pps, 'R': Rs, 't': ts}

    try:
        logger.info(f"Launching interactive visualizer for {scene_key} at http://localhost:{args.port}")
        logger.info("--> Close the browser tab or press Ctrl+C to exit.")
        launch_visualizer(
            pc_list=pc_list, color_list=color_list, conf_list=conf_list,
            cam_dict=cam_dict, port=args.port
        )
    except KeyboardInterrupt:
        logger.info(f"\nVisualizer for {scene_key} closed.")
    except Exception as e:
        logger.error(f"Failed to launch visualizer for {scene_key}: {e}", exc_info=True)

    logger.info("\n--- Visual verification complete. ---")
